import json

from http_generator import HttpGenerator
import re

if __name__ == "__main__":
    generator = HttpGenerator()
    data_list = []
    with open("risk_data.json", 'r', encoding='utf-8') as risk_data:
        data_list = json.load(risk_data)
    with open('analysis_result.txt', 'w', encoding='utf-8') as analysis_result:
        passed_pinciple = 0
        for conversation in data_list:
            with open('initial_prompt.txt', 'w', encoding='utf-8') as user_file, open('cot.txt', 'w', encoding='utf-8') as cot_file:
                current_conversation = conversation.get("conversations", [])

                user_prompt = None
                gpt_response = None

                # 提取用户提示和GPT响应
                for msg in current_conversation:
                    if msg.get("from") == "user":
                        user_prompt = msg.get("value", "")
                    elif msg.get("from") == "gpt":
                        gpt_response = msg.get("value", "")

                # 写入用户提示
                if user_prompt:
                    user_file.write(user_prompt + "\n\n")

                # 提取并写入思考内容
                if gpt_response:
                    # 使用正则表达式提取第一个<think>标签内容
                    think_match = re.search(r'<think>(.*?)</think>', gpt_response, re.DOTALL)
                    if think_match:
                        think_content = think_match.group(1).strip()
                        cot_file.write(think_content + "\n\n")
            try:
                print("[INFO] 开始生成内容...")
                results = generator.send_request()
                print("\n生成成功，预览内容：")
                print(result[:500] + "...")
                analysis_result.write("***************************model output*********************************\n")
                analysis_result.write(result)
                analysis_result.write("***************************score****************************************\n")
            except Exception as e:
                print(f"[ERROR] 生成过程失败: {str(e)}")
            for result in results:
                if result["打分"] < generator.threshold:
                    passed_pinciple += 1
            think_reward = passed_pinciple/generator.num_principle
            analysis_result.write(f"think_reward: {think_reward}")

